{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### trainer的官方api https://huggingface.co/docs/transformers/main_classes/trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 导入相关包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer # 导入trainer相关的包\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "from datasets import load_dataset\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载数据集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset = load_dataset('csv', data_files={'train': 'D:/code/datasets/ChineseNlpCorpus/datasets/ChnSentiCorp_htl_all/ChnSentiCorp_htl_all.csv'})\n",
    "# print(dataset)\n",
    "dataset = load_dataset('csv', data_files='D:/code/datasets/ChineseNlpCorpus/datasets/ChnSentiCorp_htl_all/ChnSentiCorp_htl_all.csv',split='train')\n",
    "print(dataset)\n",
    "# dataset = dataset['train'].train_test_split(test_size=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 划分数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = dataset.train_test_split(test_size=0.1)\n",
    "# 创建dataloader\n",
    "import torch\n",
    "tokenizer = AutoTokenizer.from_pretrained('D:/code/models/huggingface/rbt3')\n",
    "def process_function(examples):\n",
    "    tokenizer_examples = tokenizer(examples['review'], truncation=True, max_length=128)\n",
    "    # tokenizer_examples = tokenizer(examples['review'], truncation=True, padding='max_length', max_length=128)\n",
    "    tokenizer_examples[\"labels\"] = examples['label']\n",
    "    return tokenizer_examples\n",
    "\n",
    "tokenizer_datasets = datasets.map(process_function, batched=True,remove_columns=datasets['train'].column_names)\n",
    "len(tokenizer_datasets['test']['input_ids'][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# dataloader 处理不需要了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torch.utils.data import DataLoader \n",
    "# # 这里才正式的开始进行padding操作\n",
    "# from transformers import DataCollatorWithPadding\n",
    "train_dataset,eval_dataset  =  tokenizer_datasets['train'],tokenizer_datasets['test']\n",
    "# train_loader = DataLoader(trainset, batch_size=32, shuffle=True, collate_fn=DataCollatorWithPadding(tokenizer=tokenizer))\n",
    "# val_loader = DataLoader(valset, batch_size=64, shuffle=True, collate_fn=DataCollatorWithPadding(tokenizer=tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 优化器不需要创建了"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torch.optim import AdamW\n",
    "# optimizer =  AdamW(model.parameters(),lr=5e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 不需要创建优化器，只需要创建模型即可"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model =  AutoModelForSequenceClassification.from_pretrained('D:/code/models/huggingface/rbt3',num_labels=2)\n",
    "# cuda判定也不需要了\n",
    "# if torch.cuda.is_available():\n",
    "    # model = model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 评估部分重写evaluate的部分\n",
    "#### 1、在huggingface.co 网站中找tasks 然后在看支持的评估指标\n",
    "#### 2、直接下载失败，可以github上下载源代码，然后直接加载文件夹。\n",
    "#### 3、使用 add_batch 和 最终的compute方法最后实现评估指标的汇总工作\n",
    "#### 4、可以画出雷达图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "# accuracy = evaluate.load(\"D:/code/evaluate-main/metrics/accuracy\")\n",
    "# res = accuracy.compute(references=[0,1,0,1], predictions=[1,0,0,1])\n",
    "# print(evaluate.load('accuracy').compute(references=[1], predictions=[1]))\n",
    "# clf_metrics =  evaluate.combine(['D:/code/evaluate-main/metrics/f1','D:/code/evaluate-main/metrics/precision','D:/code/evaluate-main/metrics/recall'])\n",
    "# res\n",
    "# 4 项目都加入\n",
    "acc_metric =  evaluate.load('D:/code/evaluate-main/metrics/accuracy')\n",
    "f1_metric =  evaluate.load('D:/code/evaluate-main/metrics/f1')\n",
    "precision_metric =  evaluate.load('D:/code/evaluate-main/metrics/precision')\n",
    "recall_metric =  evaluate.load('D:/code/evaluate-main/metrics/recall')\n",
    "\n",
    "# 写一个多种标准合成的评价函数\n",
    "def eval_metrics(eval_predict):\n",
    "    predictions,labels =  eval_predict\n",
    "    predictions = predictions.argmax(axis=-1)\n",
    "    combine_acc = acc_metric.compute(predictions=predictions, references=labels)\n",
    "    f1 = f1_metric.compute(predictions=predictions, references=labels)\n",
    "    precision = precision_metric.compute(predictions=predictions, references=labels)\n",
    "    recall = recall_metric.compute(predictions=predictions, references=labels)\n",
    "    combine_acc.update(f1)\n",
    "    combine_acc.update(precision)\n",
    "    combine_acc.update(recall)\n",
    "    return combine_acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练和验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model():\n",
    "    model.eval()\n",
    "    with torch.inference_mode():\n",
    "        for batch in val_loader:\n",
    "            if torch.cuda.is_available():\n",
    "                batch = {k:v.cuda() for k,v in batch.items()}\n",
    "            output = model(**batch)\n",
    "            pred =  torch.argmax(output.logits,dim=-1)\n",
    "            # 预测batch的结果\n",
    "            clf_metrics.add_batch(predictions=pred.long(),references=batch['labels'].long())\n",
    "    # 最终拼接计算\n",
    "    return clf_metrics.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建 TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_args =  TrainingArguments(\n",
    "    output_dir='D:/code/logs/trainer/checkpoints', # 训练模型的输出路径\n",
    "    per_device_train_batch_size=64, # 训练集 batch size\n",
    "    per_device_eval_batch_size=128, # 验证集 batch size\n",
    "    \n",
    "    num_train_epochs=20, # 训练的 epochs 数\n",
    "    \n",
    "    logging_steps=150, # 每 150 个 batch 记录一次日志\n",
    "\n",
    "    # save_steps=100, # 每 100 个 batch 保存一次模型 和每个epoch 保存一次模型互斥\n",
    "    save_strategy = 'epoch', # 每 1 个 epoch 保存一次模型\n",
    "    save_total_limit=3, # 保存 1 个模型\n",
    "    \n",
    "    evaluation_strategy='epoch', # 每 1 个 epoch 验证一次模型\n",
    "\n",
    "    learning_rate=4e-5, # 学习率\n",
    "    weight_decay=0.01, # 权重衰减\n",
    "\n",
    "    metric_for_best_model='f1',\n",
    "    load_best_model_at_end=True, # 加载最优模型\n",
    "    report_to='tensorboard', # 训练日志保存到 tensorboard\n",
    "\n",
    "    )\n",
    "train_args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建  Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorWithPadding\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model, # 指定模型\n",
    "    args=train_args, # 执行训练参数\n",
    "    train_dataset=train_dataset, # 训练数据集指定 切分后的数据集\n",
    "    eval_dataset=eval_dataset, # 验证数据集指定\n",
    "    data_collator=DataCollatorWithPadding(tokenizer=tokenizer), # 数据集的预处理对数据集进行padding操作\n",
    "    compute_metrics = eval_metrics, # tokenizer\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate()\n",
    "trainer.evaluate(tokenizer_datasets['test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.predict(tokenizer_datasets['test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 保存模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def save_model(model,tokenizer,epoch):\n",
    "#     original_path = \"D:/code/models/classfication_rbt3/\"\n",
    "#     path = os.path.join(original_path,'epoch_{}'.format(epoch))\n",
    "#     if not os.path.exists(path):\n",
    "#         os.makedirs(path)\n",
    "#     model.save_pretrained(path)\n",
    "#     tokenizer.save_pretrained(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def train(epoch=5,log_step=100):\n",
    "#     global_step = 0\n",
    "#     for epoch in range(epoch):\n",
    "#         print('epoch:{}'.format(epoch))\n",
    "#         model.train()\n",
    "#         for batch in train_loader:\n",
    "#             if torch.cuda.is_available():\n",
    "#                 batch = {k:v.cuda() for k,v in batch.items()}\n",
    "#             optimizer.zero_grad()\n",
    "#             output = model(**batch)\n",
    "#             output.loss.backward()\n",
    "#             optimizer.step()\n",
    "#             if global_step % log_step == 0:\n",
    "#                 print('epoch:{},step:{},loss:{}'.format(epoch,global_step,output.loss.item()))\n",
    "#             global_step+=1  \n",
    "#         acc = evaluate_model()\n",
    "#         print('epoch:{},acc:{}'.format(epoch,acc))\n",
    "#         save_model(model,tokenizer,epoch)\n",
    "#     return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# acc = train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(type(acc),acc.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 评估结果的可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from evaluate.visualization import radar_plot\n",
    "\n",
    "test_out =  trainer.evaluate(tokenizer_datasets['test'])\n",
    "print(test_out)\n",
    "\n",
    "radar_out = [{\n",
    "        \"accuracy\": test_out['eval_accuracy'],\n",
    "        \"f1\":  test_out['eval_f1'],\n",
    "        \"precision\":  test_out['eval_precision'],\n",
    "        \"recall\":  test_out['eval_recall']\n",
    "}]\n",
    "\n",
    "model_names = [\n",
    "    \"model 1\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot =  radar_plot(data=radar_out,model_names=model_names)\n",
    "plot.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
